# -*- coding: utf-8 -*-

 #!/usr/bin/env python
import sys


import numpy as np
import argparse
import pickle as p

import os
import os.path as osp
import google.protobuf as pb
from argparse import ArgumentParser
import sys
import cv2
import caffe
from caffe.proto import caffe_pb2
from google.protobuf import text_format
from caffe import layers as LAYER
from caffe import params as PARAM

sys.path.insert(0, "xxx/caffe/python")
sys.path.append('./')
reload(sys)
sys.setdefaultencoding('utf8')

def add_mask(args):
    old_prototxt = args.src_prototxt
    old_caffemodel = args.src_caffemodel
    new_prototxt = old_prototxt.split('.prototxt')[0]+"_mask.prototxt"
    new_caffemodel =  old_caffemodel.split('.caffemodel')[0]+"_mask.caffemodel"

    # 解析src_prototxt
    old_proto= caffe_pb2.NetParameter()
    text_format.Merge(open(old_prototxt).read(), old_proto)
    new_proto= caffe_pb2.NetParameter()
    # 创建网络模型对象
    caffe.set_mode_cpu()

    top_list = []
    #写data 层
    net_old = caffe.Net(old_prototxt, old_caffemodel, caffe.TEST)


    data_layer_dim = old_proto.input_dim
    data_layer_shape = old_proto.input_shape
    data_layer_name = old_proto.input._values[0]
    if data_layer_dim != []:
        shape0 = data_layer_dim._values[0]
        shape1 = data_layer_dim._values[1]
        shape2 = data_layer_dim._values[2]
        shape3 = data_layer_dim._values[3]
    else:
        shape0 = data_layer_shape._values[0].dim._values[0]
        shape1 = data_layer_shape._values[0].dim._values[1]
        shape2 = data_layer_shape._values[0].dim._values[2]
        shape3 = data_layer_shape._values[0].dim._values[3]
    copy_layer = new_proto.layer.add()
    kwargs = {'input_param': {'shape': {'dim': [shape0,shape1,shape2,shape3]}}}
    copy_layer.CopyFrom(LAYER.Input(**kwargs).to_proto().layer[0])
    copy_layer.name = data_layer_name
    copy_layer.top[0] = data_layer_name


    #最后一层如果是softmax层或者concat层，另外处理
    special_layer = 0
    end_layer = old_proto.layer[-1]
    layer_size = len(old_proto.layer)

    if end_layer.type == "Concat":
        special_layer = 1
        layer_size = layer_size -1

    if end_layer.type == "Softmax":
        layer_size = layer_size - 1

    for i in range(layer_size):
        copy_layer = new_proto.layer.add()
        copy_layer.CopyFrom(old_proto.layer[i])

    #最后一层反卷积层输出通道数
    if special_layer == 1:
        end_layer = old_proto.layer[-1]
        if end_layer.bottom[0] in net_old.blobs:
            blob =  net_old.blobs[end_layer.bottom[0]]
            num_output = blob.channels
            top_list.append(end_layer.bottom[0])
        else:
            print("prototxt has problem!!!")
            exit()
    else:
        num_output = new_proto.layer[-1].convolution_param.num_output
        top_list.append( new_proto.layer[-1].top[0])
    #追加写mask操作层

    #增加slice层
    copy_layer = new_proto.layer.add()
    kwargs = {'slice_param': {'axis': 1}}
    copy_layer.CopyFrom(LAYER.Slice(**kwargs).to_proto().layer[0])
    copy_layer.bottom.append(top_list[0])
    copy_layer.name = "slice"
    copy_layer.top._values = []
    top_list=[]
    for i in range(num_output):
        top = "slice_ch" + "_" + str(i)
        copy_layer.top.append(top)
        top_list.append(top)

    #增加求最大值层
    copy_layer = new_proto.layer.add()
    kwargs = {'eltwise_param': {'operation': 2}}
    copy_layer.CopyFrom(LAYER.Eltwise(**kwargs).to_proto().layer[0])
    for i in range(len(top_list)):
        copy_layer.bottom.append(top_list[i])
    copy_layer.name = "eltwise_all_ch"
    copy_layer.top._values=[]
    top_list=[]
    copy_layer.top.append("slice_max")
    top_list.append("slice_max")

    last_layer = new_proto.layer[-1]
    top_list=[]
    #增加各通道减最大值层
    for i  in range(num_output):
        copy_layer = new_proto.layer.add()
        kwargs = {'eltwise_param': {'operation': 1, 'coeff': [1,-1]}}
        copy_layer.CopyFrom(LAYER.Eltwise(**kwargs).to_proto().layer[0])
        copy_layer.name = "sub_ch"+"_"+str(i)
        copy_layer.bottom.append(last_layer.bottom[i])
        copy_layer.bottom.append(last_layer.top[0])
        copy_layer.top._values=[]
        top = "slice_sub"+"_"+str(i)
        copy_layer.top.append(top)
        top_list.append(top)

    #增加个通道相加后的值concat起来
    copy_layer = new_proto.layer.add()
    kwargs = {'concat_param': {'axis': 1}}
    copy_layer.CopyFrom(LAYER.Concat(**kwargs).to_proto().layer[0])
    for i in range(len(top_list)):
        copy_layer.bottom.append(top_list[i])
    copy_layer.name = "contact_all_sub"
    copy_layer.top._values=[]
    top_list=[]
    copy_layer.top.append("contact_all_sub")
    top_list.append("contact_all_sub")

    #增加fliter层
    last_layer = new_proto.layer[-1]
    copy_layer = new_proto.layer.add()
    kwargs = {'threshold_param': {'threshold': -0.001}}
    copy_layer.CopyFrom(LAYER.Threshold(**kwargs).to_proto().layer[0])
    copy_layer.bottom.append(last_layer.top[0])
    copy_layer.name = "fliter_all_sub"
    copy_layer.top._values=[]
    top_list=[]
    copy_layer.top.append("fliter_all_sub")
    top_list.append("fliter_all_sub")

    #增加depthwsie层
    last_layer = new_proto.layer[-1]
    copy_layer = new_proto.layer.add()
    kwargs = {'convolution_param': {'num_output': num_output,'pad':0,'kernel_size':1,'stride':1,'group':num_output}}
    copy_layer.CopyFrom(LAYER.Convolution(**kwargs).to_proto().layer[0])
    copy_layer.bottom.append(last_layer.top[0])
    copy_layer.name = "depthwise"
    copy_layer.top._values=[]
    top_list=[]
    copy_layer.top.append("depthwise_all")
    top_list.append("depthwise_all")

    #增加切分层
    last_layer = new_proto.layer[-1]
    copy_layer = new_proto.layer.add()
    kwargs = {'slice_param': {'axis': 1}}
    copy_layer.CopyFrom(LAYER.Slice(**kwargs).to_proto().layer[0])
    copy_layer.bottom.append(last_layer.top[0])
    copy_layer.name = "mask_slice"
    copy_layer.top._values=[]
    top_list=[]
    for i in range(num_output):
        top = "depthwise_ch" + "_" + str(i)
        copy_layer.top.append(top)
        top_list.append(top)


    #增加mask层
    last_layer = new_proto.layer[-1]
    copy_layer = new_proto.layer.add()
    kwargs = {'eltwise_param': {'operation': 2}}
    copy_layer.CopyFrom(LAYER.Eltwise(**kwargs).to_proto().layer[0])
    for i in range(len(top_list)):
        copy_layer.bottom.append(top_list[len(top_list) - i-1])
    copy_layer.name = "mask"
    copy_layer.top._values=[]
    top_list=[]
    copy_layer.top.append("mask")
    top_list.append("mask")


    if special_layer == 1:
        end_layer = old_proto.layer[-1]
        if end_layer.bottom[1] in net_old.blobs:
            blob =  net_old.blobs[end_layer.bottom[1]]
            channels = blob.channels
            height = blob.height
            width = blob.width
            copy_layer = new_proto.layer.add()
            kwargs = {'reshape_param': {'shape':{'dim': [channels,height,width]}},'axis':1}
            copy_layer.CopyFrom(LAYER.Reshape(**kwargs).to_proto().layer[0])
            copy_layer.bottom.append(end_layer.bottom[1])
            top = end_layer.bottom[1] +"_"+ "reshape"
            copy_layer.name = top
            copy_layer.top._values=[]
            copy_layer.top.append(top)
            top_list.append(top)
        if end_layer.bottom[2] in net_old.blobs:
            blob =  net_old.blobs[end_layer.bottom[2]]
            channels = blob.channels
            height = blob.height
            width = blob.width
            copy_layer = new_proto.layer.add()
            kwargs = {'reshape_param': {'shape':{'dim': [channels,height,width]}},'axis':1}
            copy_layer.CopyFrom(LAYER.Reshape(**kwargs).to_proto().layer[0])
            top = end_layer.bottom[2] + "_" + "reshape"
            copy_layer.bottom.append(end_layer.bottom[2])
            copy_layer.name = top
            copy_layer.top._values=[]
            copy_layer.top.append(top)
            top_list.append(top)
        # concat各个通道
        copy_layer = new_proto.layer.add()
        kwargs = {'concat_param': {'axis': 1}}
        copy_layer.CopyFrom(LAYER.Concat(**kwargs).to_proto().layer[0])
        copy_layer.bottom.append(top_list[2])
        copy_layer.bottom.append(top_list[0])
        copy_layer.bottom.append(top_list[1])
        copy_layer.name = "concat"
        copy_layer.top._values=[]
        copy_layer.top.append("concat")




    with open(new_prototxt, 'w') as f:
        f.write(str(new_proto))



    net = caffe.Net(new_prototxt, caffe.TEST)

    #拷贝相同名字层的参数
    for k,v in net_old.params.items():
        # print (k,v[0].data.shape)
        # print (np.size(net_old.params[k]))
        if (k in net.params.keys()):
        # if(k in net.layer_dict.keys()):
            print(k, v[0].data.shape)
            print(np.size(net_old.params[k]))
            for i in range(np.size(net_old.params[k])):
               net.params[k][i].data[:] = np.copy(net_old.params[k][i].data[:])
    print('**********net is success*********')
    #
    #新增的通道索引层
    new_weight  = net.params['depthwise'][0].data
    # new_bias  = net.params['depthwise'][1].data
    print (np.size(new_weight))



    zeros_weight = np.zeros(new_weight.shape,float)
    # zeros_bias = np.zeros(new_bias.shape,float)


    for i in range(new_weight.shape[0]):
        # zeros_bias[i] = i+1;
        for j in range(new_weight.shape[1]):
            if special_layer == 1:
                zeros_weight[i][j] = i+1;
            else:
                zeros_weight[i][j] = i ;
                
    # for i in range(new_weight.shape[0]):
    #     for j in range(new_weight.shape[1]):
    #         if i==j:
    #            zeros_weight[i][j] = i+1;
    #            print(zeros_weight[i][j])

    new_weight[:] = np.copy(zeros_weight[:])
    # new_bias[:] = np.copy(zeros_bias[:])


    net.save(new_caffemodel)


def parse_args():
     """Parse input arguments"""
     parser = argparse.ArgumentParser(description="Generate wk_model.prototxt\n")
     parser.add_argument('--src_prototxt', '-sp', dest='src_prototxt',
                         help='Path to source prototxt', type=str)
     parser.add_argument('--src_caffemodel', '-sc', dest='src_caffemodel',
                         help='Path to source caffemodel', type=str)

     if len(sys.argv) == 1:
         parser.print_help()
         sys.exit(1)

     args = parser.parse_args()
     return args

if __name__ == "__main__":
    args = parse_args()
    add_mask(args)